/*
 * Name: elfswitch
 * Version: 0.1.0
 * Description: Very simple program to make
 * an ELF file momentarily unexecutable by 
 * modifying the entry address.
 * (Apply again to nullify the effect)
 * This program can be useful when
 * handling infected executables.
 * Dependencies: bzbutil, bzbfile
*/

#include "bzbfile.h"
#include "bzbio.h"
#include "bzbstring.h"

using namespace bzbio;
using namespace bzbstring;


constexpr int VERSION_MAJOR = 0;
constexpr int VERSION_MINOR = 1;
constexpr int VERSION_PATCH = 0;

const char* PROGNAME = "elfswitch";


void print_help_exit(int err = 0) {
	println("Usage: ", PROGNAME, " [ -h | --help ] | <filename>");
	println("Negate the entry address of the specified ELF to",
			" make it unexecutable");
	println("-h | --help\tShow this help");
	exit(err);
}


void print_version_exit() {
	print(PROGNAME, " v");
	println(VERSION_MAJOR, ".", VERSION_MINOR, ".", VERSION_PATCH);
	println("This program is part of the Balzebub project.");
	exit(0);
}


int main(int argc, char const *argv[]) {

	if(argc != 2) {
		print_help_exit(1);
	}

	std::string arg(argv[1]);

	if(arg == "-v" || arg == "--version") {
		print_version_exit();
	}

	if(arg == "-h" || arg == "--help") {
		print_help_exit();
	}

	std::string filename = arg;
	bool is_32bit;

	try {
		is_32bit = bzbfile::get_elf_class(filename) == bzbfile::ELF_CLASS_32;
	} catch(...) {
		exit_on_error("ERROR Unable to read file", 2);
	}

	std::vector<char> data;

	try {
		if(is_32bit) {
			data = bzbfile::read_bin(filename, 4, 24);
		} else {
			data = bzbfile::read_bin(filename, 8, 24);
		}
	} catch(...) {
		exit_on_error("ERROR Unable to read file", 2);
	}

	if(data.size() != 8 && data.size() != 4) {
		exit_on_error("ERROR The file is corrupted", 2);
	}
	
	uint64_t address = 0x0;
	readint(data, 0, &address);

	println("Old entry address: 0x", std::hex, address, std::dec);
	println("New entry address: 0x", std::hex, ~address, std::dec);

	for(size_t i = 0; i < data.size(); ++i) {
		uint8_t c = ~(*((uint8_t*) &data[i]));
		data[i] = *((char*) &c);
	}

	try {
		bzbfile::write_bin(filename, 24, data);
	} catch(...) {
		exit_on_error("ERROR Unable to write to file", 3);
	}

	return 0;
}
